import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import LinearNDInterpolator
from astropy.modeling import models, fitting

from scipy.ndimage import gaussian_filter
from matplotlib.colors import LogNorm
import sys

def cart2pol(x,y):
    rho = np.sqrt(x**2 + y ** 2)
    phi = np.arctan2(y,x)
    return(rho,phi)

def pol2cart(rho,phi):
    x = rho * np.cos(phi)
    y = rho * np.sin(phi)
    return(x,y)

# from numpy.core.umath_tests import matrix_multiply


def spherical_to_cartesian(latit, longit,radius):
    
    theta=np.deg2rad(latit+90)
    phi=np.deg2rad(longit)
    
    
    x = radius*np.cos(phi) * np.sin(theta)
    y = radius* np.sin(phi) * np.sin(theta)
    z = radius* np.cos(theta)
    return x, y, z

def cartesian_to_spherical(x, y, z):
    radius = np.sqrt(x ** 2 + y ** 2 + z ** 2)
    theta = np.arctan2(np.sqrt(x ** 2 + y ** 2), z)
    phi = np.arctan2(y, x)
    if x.any() < 0: phi[x < 0] = np.arctan2(y[x < 0], x[x < 0]) + np.pi
    
    theta_d=np.rad2deg(theta)
    phi_d=np.rad2deg(phi)

    phi_d[phi_d < 0]=phi_d[phi_d < 0]+360
    theta_d=theta_d-90

    return theta_d, phi_d, radius

def unit_vector(data, axis=None, out=None):

    if out is None:
        data = np.array(data, dtype=np.float64, copy=True)
        if data.ndim == 1:
            data /= np.sqrt(np.dot(data, data))
            return data
    else:
        if out is not data:
            out[:] = np.array(data, copy=False)
        data = out
    length = np.atleast_1d(np.sum(data*data, axis))
    np.sqrt(length, length)
    if axis is not None:
        length = np.expand_dims(length, axis)
    data /= length
    if out is None:
        return data


def rotation_matrix(angle, direction, point=None):

    sina = np.sin(angle)
    cosa = np.cos(angle)
    direction = unit_vector(direction[:3])
    # rotation matrix around unit vector
    R = np.diag([cosa, cosa, cosa])
    R += np.outer(direction, direction) * (1.0 - cosa)
    direction *= sina
    R += np.array([[ 0.0,         -direction[2],  direction[1]],
                      [ direction[2], 0.0,          -direction[0]],
                      [-direction[1], direction[0],  0.0]])
    M = np.identity(4)
    M[:3, :3] = R
    if point is not None:
        # rotation not around origin
        point = np.array(point[:3], dtype=np.float64, copy=False)
        M[:3, 3] = point - np.dot(R, point)
    return M


def rotatePole(lats, lons, radius, angle=90, axis=[1,0,0]):
    """
    Rotates the given geodetic lat/lon coordinates around the origin.
    
    :param lats, lons: shape (n,) in radians
    :param altitude: in km
    :param angle: degrees
    :param axis: [1, 0, 0], [0, 1, 0], or [0, 0, 1] for x y z axis
    :rtype: tuple (lats, lons) in radians
    """
    assert lats.size == 1 and lons.size == 1
    assert len(axis) == 3    

    x,y,z = spherical_to_cartesian(lats, lons, radius)
    xyz = np.asarray([x,y,z]).T
    
    print(xyz.shape)
    
    alpha = np.deg2rad(angle)
    rot = rotation_matrix(alpha, axis)[:3,:3]
    
    # xyzRot = matrix_multiply(rot,xyz[...,np.newaxis]).reshape(xyz.shape)
    xyzRot = np.matmul(rot,xyz[...,np.newaxis]).reshape(xyz.shape)
    
    lats, lons, radius = cartesian_to_spherical(xyzRot[:,0], xyzRot[:,1], xyzRot[:,2])
    

    
    return lats, lons, radius




# 01 aug 2017
S_sel = 31.8
S_diameter = 17.8


# %%
##### load and set up the actual lat and lon values

nc_lat = np.load('lat.npy')

nc_lat = np.flip(nc_lat[152:272,50:250])

plt.imshow(nc_lat)

nc_lon = np.load('lon.npy')

nc_lon = (360-np.flip(nc_lon[152:272,50:250]))/360*24

fake_lon = nc_lon + 0
fake_lon[fake_lon > 12] = 24 - fake_lon[fake_lon > 12]


idx = np.unravel_index(np.argmax(nc_lat), nc_lat.shape)

pole_pos_x = idx[1]
pole_pos_y = idx[0]

#  size of equator in 0.1" pixels

# 01 aug 2017
S_sel = 31.8
S_diameter = 17.8


nc_lat[nc_lon == 24.0] = np.nan
nc_lat[nc_lon == 18.0] = np.nan

nc_lon[nc_lon == 24.0] = np.nan
nc_lon[nc_lon == 18.0] = np.nan

# plt.imshow(nc_lon,origin='lower')

# plt.figure()
# plt.plot(nc_lon[100,:])
# print(nc_lon[100,70])



# %%
##### set up the radian velocity on the planet, caused by corotation - this is our radial scaler, but not yet radial



scaling_vel = (9.87)#/(S_diameter*10/2)

rotvel=np.cos(np.radians(nc_lat))*scaling_vel

# plt.figure()

# plt.plot(rotvel[60,:])


# plt.plot(rotvel)



pole_lat = nc_lat
pole_lon = nc_lon




# %%
##### convert the corot_scaler into line of sight from Earth (excluding disk angle) i.e. vector in Earth direction around 

### radial velocities are at a maximum at 90 and 270, and at a minimum at 0 and 180 (the reverse of corotation)


nc_lon_rad = np.radians(nc_lon/24.*360)

phase_component = (np.cos(nc_lon_rad))

# plt.figure()

# plt.imshow(phase_component,vmin=-1,vmax=1)



# %%
##### add in the disk position effect, as at the sub-solar point, you see no component of this velocity 
### radial velocities are at a maximum at 90 and 270, and at a minimum at 0 and 180 (the reverse of corotation)

# 1) calculate the emission angle -  so squish the planet then comapre with the centre of the planet

# yvalue=1.5

# ydepth = yvalue*2
# pix_ydepth = nc_lat[:,0].size

# # planet flattening

# orbitalinc=np.cos(np.radians(S_sel))
# flattening=1 - (0.0979620*orbitalinc)


# #  mathematically, it is better to make the planet into a circle and calculate the values from there
# unflat_depth= ydepth/flattening



# unflat_Yheight = S_diameter - unflat_depth

# plt.figure()
# fig, ax = plt.subplots(2, figsize=(5*1.,3.2*1.4))

# velplot=ax[0].imshow(phase_component,vmin=-2.0,vmax=2.0,cmap='jet_r',origin='lower',extent=[-10,+10,-yvalue,+yvalue])


# Xdistance = (np.arange(nc_lat[0,:].size)-100)*0.1

# unflat_Ydistance = (unflat_Yheight + np.arange(pix_ydepth)/pix_ydepth*unflat_depth)

# x=100
# y=60


# x = np.linspace(0,199,200)
# y = np.linspace(0,119,120)
# x_1, y_1 = np.meshgrid(x,y)

# Xdistance = (x_1-100)*0.1
# unflat_Ydistance= (unflat_Yheight + y_1/pix_ydepth*unflat_depth)


# unflat_distance = np.sqrt(Xdistance**2 + unflat_Ydistance**2)

# emission_angle_ratio = unflat_distance/S_diameter

# emission_angle_ratio[emission_angle_ratio > 1] = np.nan


# velplot2=ax[1].imshow(emission_angle_ratio,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=0,vmax=1)


# print(x_1[0,1],y_1[0,1],Xdistance[0,1],unflat_Ydistance[0,1],unflat_distance[0,1],emission_angle_ratio[0,1])


# that was messy and ended up not really working - lets try just convering the lat long to looking down on the pole and 90-lat will give you the values I think

# degree_lon = nc_lon/24*360

# emission_angle = np.zeros_like(degree_lon)

# # need to loop through x and y for this to work
# for j in range(degree_lon[:,0].size):
#     for i in range(degree_lon[0,:].size):

#         lat0_shift1, lon0_shift1, radius = rotatePole(nc_lat[j,i], degree_lon[j,i],1,angle=S_sel-90, axis=[0,1,0])
#         emission_angle[j,i] = lat0_shift1



# # FUCK PYTHON - holy shit this is so fucking shitty - numpy.core.umath_testsis fucked, so back to square one with development AGAIN
# # Nothing in this stupid program ever works after a year, its just fucking weak

# plt.figure()
# plt.imshow(emission_angle)








# %% ### use Sorba to map from ionosphere to magnetosphere R across local times 


South = "False"
Storm = "True"


if Storm == "True":
    storm_str='_comp'
else:
    storm_str='_expand'


if South == "True":
    noon_filename='Snoon_L_colat'+storm_str+'.txt'
    dawn_filename='Sdawn_L_colat'+storm_str+'.txt'
    dusk_filename='Sdusk_L_colat'+storm_str+'.txt'
    night_filename='Snight_L_colat'+storm_str+'.txt'
    lat_min = 58
    south_str='_south'
else:
    noon_filename='Nnoon_L_colat'+storm_str+'.txt'
    dawn_filename='Ndawn_L_colat'+storm_str+'.txt'
    dusk_filename='Ndusk_L_colat'+storm_str+'.txt'
    night_filename='Nnight_L_colat'+storm_str+'.txt'
    lat_min=62.6
    south_str='_north'


# noon_filename='Nnoon_L_colat.txt'
noon=np.loadtxt(noon_filename)

# dawn_filename='Ndawn_L_colat.txt'
dawn=np.loadtxt(dawn_filename)

# dusk_filename='Ndusk_L_colat.txt'
dusk=np.loadtxt(dusk_filename)

# night_filename='Nnight_L_colat.txt'
night=np.loadtxt(night_filename)


noon_filename='Wilson_Noon_radial.csv'
noonvel=np.loadtxt(noon_filename)

dawn_filename='Wilson_Dawn_radial.csv'
dawnvel=np.loadtxt(dawn_filename)

dusk_filename='Wilson_Dusk_radial.csv'
duskvel=np.loadtxt(dusk_filename)

night_filename='Wilson_Midnight_radial.csv'
nightvel=np.loadtxt(night_filename)






noon_lat=noon[:,1]
noon_long=np.zeros_like(noon_lat)+12
noon_R=noon[:,0]


dawn_lat=dawn[:,1]
dawn_long=np.zeros_like(dawn_lat)+6
dawn_R=dawn[:,0]

night_lat=night[:,1]
night_long=np.zeros_like(night_lat)
night_R=night[:,0]

dusk_lat=dusk[:,1]
dusk_long=np.zeros_like(dusk_lat)+18
dusk_R=dusk[:,0]


x = np.hstack([night_long,0,dawn_long,6,noon_long,12,dusk_long,18,night_long+24,24])
y = np.hstack([night_lat,90,dawn_lat,90,noon_lat,90,dusk_lat,90,night_lat,90])
z = np.hstack([night_R,1,dawn_R,1,noon_R,1,dusk_R,1,night_R,1])



X = np.linspace(min(x), max(x))
Y = np.linspace(min(y), max(y))

X2 = pole_lon
Y2= 90-pole_lat


X, Y = np.meshgrid(X, Y)  # 2D grid for interpolation

interp = LinearNDInterpolator(list(zip(x, y)), z)


Z = interp(X2, Y2)

Zplot = Z+0
Zplot[nc_lat < 1] = np.nan

yvalue=1.5

levels = np.arange(0,90,10)
levels2 = np.arange(0,24,2)



yvalue=1.5

levels = np.arange(0,90,10)
levels2 = np.arange(0,24,2)

# ax[0].pcolormesh(X2, Y2, Z, shading='auto')
# plot_mapping=ax[0].imshow(Zplot,origin='lower',extent=[-10,+10,-yvalue,+yvalue], norm=LogNorm(vmin=1, vmax=50))
# ax[0].contour(nc_lat,levels,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[0].contour(fake_lon,levels2,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[0].plot([0,0],[-1,1.49],color='grey',linestyle='dotted')
# mcb=plt.colorbar(plot_mapping,ax=ax[0],label='Magnetospheric equatorial R$_S$', orientation="horizontal", ticks=[1, 10, 50])
# mcb.ax.set_xticklabels(['1', '10', '50']) 

# ax[0].set_xlim([-7,7])



# %%



veldiff_model=np.load('model_vels'+south_str+storm_str+'.npy')
veldiff_model_r15=np.load('model_vels'+south_str+storm_str+'_r15.npy')
veldiff_model_r6=np.load('model_vels'+south_str+storm_str+'_r6.npy')
Zplot_model=np.load('model_dist'+south_str+storm_str+'.npy')


# %%

#######  Use wilson to plot from R to a subrotation rate 

wilson_filename='wilson2.txt'
wilson=np.loadtxt(wilson_filename)


# calculate the corotation speed

# 40 x 4s and 40 x 30 making the start and end limits very hard to avoid for the polynomial 
x = np.hstack([4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,wilson[:,0],31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31])
y = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,wilson[:,1]/(wilson[:,0]*9.87),0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333])

xx=np.arange(34)+1

g_init = models.Polynomial1D(degree=4) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
fit_g = fitting.LevMarLSQFitter()
g_wilson = fit_g(g_init, x, y,maxiter=1000)

g_wilson_xx = g_wilson(xx)
g_wilson_xx[xx<4]=1
g_wilson_xx[xx>30]=0.333


# %%


#########  Figure out the Wilson radial velocities

noonvel_R = noonvel[:,0]

noonvel_co = noonvel[:,1]/(noonvel[:,0]*9.87)

dawnvel_R = dawnvel[:,0]
dawnvel_co = dawnvel[:,1]/(dawnvel[:,0]*9.87)

nightvel_R = nightvel[:,0]
nightvel_co = nightvel[:,1]/(nightvel[:,0]*9.87)

duskvel_R = duskvel[:,0]
duskvel_co = duskvel[:,1]/(duskvel[:,0]*9.87)




fig, axes = plt.subplots(nrows=2, ncols=2)

fig.set_size_inches(8, 5)
                    
ax = axes.ravel()


ax[0].set_xlabel('Magnetospheric equatorial R$_S$')


ax[0].set_ylabel('Azimuthal velocity [km/s]')
ax[1].set_xlabel('Magnetospheric equatorial R$_S$')
ax[1].set_ylabel('Sub-corotation')

ax[2].set_xlabel('Magnetospheric equatorial R$_S$')
ax[2].set_ylabel('$\Delta$ Sub-corotation')


ax[3].set_xlabel('Magnetospheric equatorial R$_S$')
ax[3].set_ylabel('Modelled Sub-corotation')

# ax[1].plot(noonvel_R, noonvel_co, marker=".",linestyle='None',c='blue')
# ax[1].plot(dawnvel_R, dawnvel_co, marker=".",linestyle='None',c='pink')
# ax[1].plot(nightvel_R, nightvel_co, marker=".",linestyle='None',c='k')
# ax[1].plot(duskvel_R, duskvel_co, marker=".",linestyle='None',c='green')


# ax[1].set_xlabel('Magnetospheric equatorial R$_S$')
# ax[1].set_ylabel('Sub-corotation')

xx=np.arange(34)+1


noonx = np.hstack([1,2,3,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,noonvel_R,30,31,32])
noony = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,noonvel_co,0.144,0.144,0.144])

# noonx = np.hstack([4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,noonvel_R,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31])
# noony = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,noonvel_co,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333])


noong_init = models.Polynomial1D(degree=10) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
noonfit_g = fitting.LevMarLSQFitter()
noong = noonfit_g(noong_init, noonx, noony,maxiter=1000)

noongxx = noong(xx)
noongxx[xx<4]=1
noongxx[xx>28]=g_wilson_xx[xx > 28]

print(noong.param_names)


    # Plot the data with the best-fit model
    # plt.figure(figsize=(8,5))
# ax[2].plot(noonx, noony, marker=".",linestyle='None',c='blue')
# ax[2].plot(xx, noongxx,c='darkblue')
ax[0].plot(noonvel[:,0], noonvel[:,1],c='darkblue',marker='.',linestyle='None')
ax[1].plot(xx, noongxx,c='darkblue')
ax[1].plot(noonx, noony,c='darkblue',linestyle='none',marker=',')
ax[2].plot(xx, noongxx-g_wilson_xx,c='darkblue')
# ax[3].set_xlabel('Magnetospheric equatorial R$_S$')
# ax[3].set_ylabel('Sub-corotation')



dawnx = np.hstack([4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,dawnvel_R,29,30,31])
dawny = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,dawnvel_co,0.146001,0.146001,0.146001])

dawng_init = models.Polynomial1D(degree=7) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
dawnfit_g = fitting.LevMarLSQFitter()
dawng = dawnfit_g(dawng_init, dawnx, dawny,maxiter=1000)

dawngxx = dawng(xx)
dawngxx[xx<4]=1
dawngxx[xx>15]=g_wilson_xx[xx > 15]

print(dawng.param_names)


    # Plot the data with the best-fit model
    # plt.figure(figsize=(8,5))
# ax[2].plot(dawnx, dawny, marker=".",linestyle='None',c='lightpink')
# ax[2].plot(xx, dawngxx,c='pink')

ax[0].plot(dawnvel[:,0], dawnvel[:,1],c='pink',marker='.',linestyle='None')
ax[1].plot(dawnx, dawny,c='pink',linestyle='none',marker=',')
ax[1].plot(xx, dawngxx,c='pink')
ax[2].plot(xx, dawngxx-g_wilson_xx,c='pink')




nightx = np.hstack([4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,nightvel_R,29,30,31])
nighty = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,nightvel_co,0.146001,0.146001,0.146001])

nightg_init = models.Polynomial1D(degree=4) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
nightfit_g = fitting.LevMarLSQFitter()
nightg = nightfit_g(nightg_init, nightx, nighty,maxiter=1000)

nightgxx = nightg(xx)
nightgxx[xx<4]=1
nightgxx[xx>28]=g_wilson_xx[xx > 28]

print(nightg.param_names)


    # Plot the data with the best-fit model
    # plt.figure(figsize=(8,5))
# ax[2].plot(nightx, nighty, marker=".",linestyle='None',c='grey')
# ax[2].plot(xx, nightgxx,c='black')

corot_pos=np.arange(5,31)

ax[0].plot(nightvel[:,0], nightvel[:,1],c='black',marker='.',linestyle='None')
ax[1].plot(xx, nightgxx,c='black')
ax[1].plot(nightx, nighty,c='black',linestyle='none',marker=',')
ax[0].plot(corot_pos,corot_pos *9.87,c='red',linestyle='dashed')
ax[0].plot(xx, g_wilson_xx*xx*9.87,c='orange',linestyle='dashed')
ax[1].plot(xx, g_wilson_xx,c='orange',linestyle='dashed')
ax[2].plot(xx, nightgxx-g_wilson_xx,c='black')
# ax[3].set_xlabel('Magnetospheric equatorial R$_S$')
# ax[3].set_ylabel('Sub-corotation')


duskx = np.hstack([4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,duskvel_R,29,30,31])
dusky = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,duskvel_co,0.146001,0.146001,0.146001])

duskg_init = models.Polynomial1D(degree=9) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
duskfit_g = fitting.LevMarLSQFitter()
duskg = duskfit_g(duskg_init, duskx, dusky,maxiter=1000)

duskgxx = duskg(xx)
duskgxx[xx<4]=1
duskgxx[xx>28]=g_wilson_xx[xx > 28]

print(duskg.param_names)


    # Plot the data with the best-fit model
    # plt.figure(figsize=(8,5))
# ax[2].plot(duskx, dusky, marker=".",linestyle='None',c='green')
# ax[2].plot(xx, duskgxx,c='darkgreen')
ax[0].scatter(duskvel[:,0], duskvel[:,1],c='darkgreen',marker='.')
ax[1].plot(xx, duskgxx,c='darkgreen')
ax[1].plot(duskx, dusky,c='darkgreen',linestyle='none',marker=',')
ax[2].plot(xx, duskgxx-g_wilson_xx,c='darkgreen')
# ax[3].set_xlabel('Magnetospheric equatorial R$_S$')
# ax[3].set_ylabel('Sub-corotation')



# ax[3].plot(xx[:15], dawngxx[:15]-duskgxx[:15],c='maroon')
# ax[3].plot(xx[15:], dawngxx[15:]-duskgxx[15:],c='maroon',linestyle='dotted')
# ax[3].plot(xx, nightgxx-noongxx,c='purple')



noon_R=xx
noon_LT=np.zeros_like(noon_R)+12
noon_Rv=noongxx


dawn_R=xx
dawn_LT=np.zeros_like(dawn_R)+6
dawn_Rv=dawngxx


# ###### make more smooth and normalised between DuNi and DaNo


night_R=xx
night_LT=np.zeros_like(night_R)
night_Rv=nightgxx

dusk_R=xx
dusk_LT=np.zeros_like(dusk_R)+18
dusk_Rv=duskgxx

# ###### make more smooth and normalised between DuNi and DaNo

dusk_Rv=nightgxx






yyy = np.hstack([night_LT,dawn_LT,noon_LT,dusk_LT,night_LT+24])
xxx = np.hstack([night_R,dawn_R,noon_R,dusk_R,night_R])
zzz = np.hstack([night_Rv,dawn_Rv,noon_Rv,dusk_Rv,night_Rv])


# X = np.linspace(min(x), max(x))
# Y = np.linspace(min(y), max(y))

# X2 = pole_lon
# Y2= 90-pole_lat


# X, Y = np.meshgrid(X, Y)  # 2D grid for interpolation

interp = LinearNDInterpolator(list(zip(xxx, yyy)), zzz)



rr11 = np.linspace(0,40,41)
lt11 = np.linspace(0,24,25)
X2, Y2 = np.meshgrid(rr11,lt11)


ZZ = interp(X2, Y2)

ZZplot = ZZ+0
# ZZplot[nc_lat < 1] = np.nan

yvalue=1.5

levels = np.arange(0,90,10)
levels2 = np.arange(0,24,2)

# ax[0].pcolormesh(X2, Y2, Z, shading='auto')
# plot_mapping=ax[3].imshow(ZZplot,origin='lower')#,extent=[-10,+10,-yvalue,+yvalue], norm=LogNorm(vmin=1, vmax=50))
# ax[3].contour(nc_lat,levels,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[3].contour(fake_lon,levels2,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[3].plot([0,0],[-1,1.49],color='grey',linestyle='dotted')
# mcb=plt.colorbar(plot_mapping,ax=ax[0],label='Magnetospheric equatorial R$_S$', orientation="horizontal", ticks=[1, 10, 50])
# mcb.ax.set_xticklabels(['1', '10', '50']) 

# ax[0].set_xlim([-7,7])


# %% this has modelled the velocity with R and LT, so now we need to pipe that back into the map


# Zplot is the map of distance



# fig, axes = plt.subplots(nrows=3, ncols=1)

# fig.set_size_inches(8, 5)
                    
# ax = axes.ravel()



# ax[0].pcolormesh(X2, Y2, Z, shading='auto')
# plot_mapping=ax[0].imshow(Zplot,origin='lower',extent=[-10,+10,-yvalue,+yvalue], norm=LogNorm(vmin=1, vmax=50))
# ax[0].contour(nc_lat,levels,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[0].contour(fake_lon,levels2,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[0].plot([0,0],[-1,1.49],color='grey',linestyle='dotted')
# mcb=plt.colorbar(plot_mapping,ax=ax[0],label='Magnetospheric equatorial R$_S$', orientation="horizontal", ticks=[1, 10, 50])
# mcb.ax.set_xticklabels(['1', '10', '50']) 

# ax[0].set_xlim([-7,7])


# plot_LT = ax[1].imshow(nc_lon,origin='lower',extent=[-10,+10,-yvalue,+yvalue])
# mcb=plt.colorbar(plot_LT,ax=ax[1],label='Localtime', orientation="horizontal")
# mcb.ax.set_xticklabels(['1', '10', '50']) 


# so LT is nc_lon and distance is Zplot:
    
radial_flow = interp(Zplot, nc_lon)


# plot_radial = ax[2].imshow(radial_flow,origin='lower',extent=[-10,+10,-yvalue,+yvalue])
# mcb=plt.colorbar(plot_radial,ax=ax[2],label='Localtime', orientation="horizontal")






    # Plot the data with the best-fit model
    # plt.figure(figsize=(8,5))
# ax[2].plot(x, y, marker=".",linestyle='None')
# ax[2].plot(xx, gxx,c='darkred')
# ax[2].set_xlabel('Magnetospheric equatorial R$_S$')
# ax[2].set_ylabel('Sub-corotation')

# veldiff_model=np.load('model_vels.npy')
# Zplot_model

ax[3].plot(Zplot_model[90:,90], veldiff_model[90:,90]*(-1),c='black')
ax[3].plot(Zplot_model[0:90,90], veldiff_model[0:90,90]*(-1),c='darkblue')
ax[3].plot(Zplot_model[90,90:], veldiff_model[90,90:]*(-1),c='darkgreen')
ax[3].plot(Zplot_model[90,:90], veldiff_model[90,:90]*(-1),c='pink')


bbox_props2_1 = dict(boxstyle="circle,pad=0.05", fc="lightcyan", ec="dodgerblue", lw=2)

ax[0].text(1,260,'a',bbox=bbox_props2_1)
ax[1].text(1,0.9,'b',bbox=bbox_props2_1)
ax[2].text(1,0.12,'c',bbox=bbox_props2_1)
ax[3].text(1,0.14,'d',bbox=bbox_props2_1)


# ax[3].plot(Zplot_model[90:,90], veldiff_model_r15[90:,90]*(-1),c='black',linestyle='dashed')
# ax[3].plot(Zplot_model[0:90,90], veldiff_model_r15[0:90,90]*(-1),c='darkblue',linestyle='dashed')
# ax[3].plot(Zplot_model[90,90:], veldiff_model_r15[90,90:]*(-1),c='darkgreen',linestyle='dashed')
# ax[3].plot(Zplot_model[90,:90], veldiff_model_r15[90,:90]*(-1),c='pink',linestyle='dashed')



# ax[3].xlabel('Position')
# ax[3].ylabel('Flux')












# %%

# OK - next, we need to read in the values from the Wilson paper and convert then into corotation (*not* subcorot, but corot) values, to produce a 0-1 factor for the above and map it into the planet
emission_angle_ratio = 1 #(until it is fixed)

vscale = rotvel*phase_component*emission_angle_ratio*radial_flow

# The above is right - the radial flow is towards you at the bottom *in the magnetosphere* but is away from you in the pole - because the direction flips with the mapping of the magnetic field line!!!

vscale[np.isnan(vscale)]=0

# plt.figure()
# fig, ax = plt.subplots(4, figsize=(5*1.,3.2*1.4))

yvalue=1.5

# velplot=ax[0].imshow(rotvel,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=0,vmax=9)

# velplot2=ax[1].imshow(phase_component,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=-1,vmax=1)

# velplot3=ax[1].imshow(radial_flow,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=-0.5,vmax=0.5,cmap='seismic_r')
# mcb=plt.colorbar(velplot3,ax=ax[1],label='Radial flow (% corotation)', orientation="horizontal")



# velplot4=ax[3].imshow(vscale,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=-0.5,vmax=0.5,cmap='seismic_r')
# mcb=plt.colorbar(velplot4,ax=ax[3],label='Radial vel component', orientation="horizontal")

# ax[3].set_xlabel('arcsec (x)')
# ax[3].set_ylabel('arcsec (y)')
# # ax[3].contour(Xplot,[0.5,1.5],colors='orange',extent=[-10,+10,-yvalue,+yvalue])
# # ax[3].contour(Zplot,[0,18,100],colors='blue',extent=[-10,+10,-yvalue,+yvalue])

# ax[3].set_ylim([-1.5,1.5])
# ax[3].set_xlim([-4,4])

fig.tight_layout(pad=1.2)
plt.savefig('asd_azimuthal_vel.pdf', dpi=300, bbox_inches='tight', facecolor='white', pad_inches=0) 
plt.show()

np.save('azimuthal_velocity'+south_str+storm_str+'.npy',vscale)



# %% now we combine the above into a velocity plot



#  for now, without the corot mapped values




# %%

# ####  set up a scaling set of latitudes across the array


# ##### a quick look at the data suggest 34 pixels between  the main emission at ~12 colatitude  either side of the pole

# # scaling_width = 80
# scaling_width = 90

# #84 = 36 pixels apart, needs to be 34
# #80 = 34 pixels apart, same as auroral emission at 12 colatitude

# pole_lat = np.zeros([scaling_width*2+1,scaling_width*2+1])


# rotvel=(np.arange(scaling_width*2+1)-scaling_width)/scaling_width*(-9.87)
# plt.figure()
# plt.plot(rotvel)

# for ii in range (scaling_width*2+1):
#     for jj in range (scaling_width*2+1):
#         i = scaling_width - ii
#         j = scaling_width - jj
#         if np.sqrt(i**2 + j**2) < scaling_width:
#             ratio = np.sqrt(i**2 + j**2) / scaling_width
#             pole_lat[jj,ii]=np.degrees(np.arccos(ratio))
        

# pole_lat2=pole_lat+0
# pole_lat2[pole_lat2 < 78]=0
# pole_lat2[pole_lat < 77]=pole_lat[pole_lat < 77]

# pl2=pole_lat2[int(scaling_width*0.75):int(scaling_width*1.25),int(scaling_width*0.75):int(scaling_width*1.25)]
# pl3=pole_lat2[scaling_width,int(scaling_width*0.75):int(scaling_width*1.25)]


# # plt.figure()
# # # plt.imshow(pl2)
# # plt.plot(pole_lat2[scaling_width,int(scaling_width*0.75):int(scaling_width*1.25)])


# ####  set up a scaling set of longitudes across the array



# pole_lon = np.zeros([scaling_width*2+1,scaling_width*2+1])

# for ii in range (scaling_width*2+1):
#     for jj in range (scaling_width*2+1):
#         i = scaling_width - ii
#         j = scaling_width - jj
#         rho,phi = cart2pol(i,j)

#         if np.sqrt(i**2 + j**2) < scaling_width:
#             pole_lon[jj,ii]= (360-np.mod((np.degrees(phi)+270) ,360))/360*24
        




# # pole_lon2=pole_lon+0
# # pole_lon2[pole_lon2 < 76]=0
# # pole_lon2[pole_lon < 74]=pole_lon[pole_lon < 74]








# cov = fit_g.fit_info['param_cov']


# ax[3].set_title("Downscaled image (no aliasing)")

# ax[0].set_xlim(0, 512)
# ax[0].set_ylim(512, 0)








# #######  Map from Sorba to Wilson to map subrotation rate across the planet 

# # firstly, if mapping is Nan, that means you are in the polar cap, so use a subcorotation of 0.3 (or some other escape clause)



# subrot=g(Z)

# subrot[Z < 4]=1
# subrot[Z > 30]=0.3333

# # subrot[np.isnan(subrot)]=0.3333
# subrot[np.isnan(subrot)]=0.3

# losvel=np.zeros_like(subrot)

# subrotplot = subrot+0
# subrotplot[nc_lat < 1] = np.nan

# # ax[4].plot(Z)
# plot_subrot=ax[1].imshow(subrotplot,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=0,vmax=1)
# # ax[5].set_ylim([0,15])
# ax[1].contour(nc_lat,levels,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[1].contour(fake_lon,levels2,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax[1].plot([0,0],[-1,1.49],color='grey',linestyle='dotted')
# plt.colorbar(plot_subrot,ax=ax[1],label='Sub-corotation', orientation="horizontal")
# ax[1].set_xlim([-7,7])


# for i in range(subrot[:,0].size):
#     losvel[i,:]=subrot[i,:]*rotvel

# seeing=0.67
# losvelseeing = gaussian_filter(losvel, sigma=seeing*10/2.355)

# # 0.5" is 5 pixel - but is FWHM - convert to sigma fwhm = 2.355 sigma

 
# # ax[5].imshow(losvel)
# arcsec_x = np.arange(200)*0.1-10

# ax[3].plot(arcsec_x,losvel[pole_pos_y,:])
# ax[3].plot(arcsec_x,losvelseeing[pole_pos_y,:])
# ax[3].set_xlim([-7,7])
# ax[3].set_ylim([-5,5])
# ax[3].set_xlabel('arcsec (x)')
# ax[3].set_ylabel('Velocity (km/s)')

# plt.tight_layout()


# np.save('losvel.npy',losvel)
# np.save('losvelseeing.npy',losvelseeing)


# plt.savefig('asd_corot_model.pdf', dpi=300, bbox_inches='tight', facecolor='white', pad_inches=0) 
# plt.show()
